import random
import numpy as np
from torchvision.ops import box_iou as iou
import cv2
from torch.nn import Module, Sequential
import torch


class SwapChannels(Module):
    def __init__(self, swaps):
        super().__init__()
        self.swaps = swaps

    def forward(self, image):
        image = image[:, :, self.swaps]
        return image


class RandomLightingNoise(Module):
    def __init__(self):
        super().__init__()
        self.perms = ((0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0))

    def forward(self, image):
        if random.randint(0, 1):
            swap = random.choice(self.perms)
            shuffle = SwapChannels(swap)
            image = shuffle(image)
        return image


class RandomBrightness(Module):
    def __init__(self, delta=32):
        super().__init__()
        self.delta = delta

    def forward(self, image):
        if random.randint(0, 1):
            delta = random.uniform(-self.delta, self.delta)
            # 图像中的每个像素加上一个随机值
            image += delta
        return image


class RandomHue(Module):
    def __init__(self, delta=18.0):
        super().__init__()
        self.delta = delta

    def forward(self, image):
        if random.randint(0, 1):
            image[:, :, 0] += random.uniform(-self.delta, self.delta)
            # 规范超过范围的像素值
            image[:, :, 0][image[:, :, 0] > 360.0] -= 360.0
            image[:, :, 0][image[:, :, 0] < 0.0] += 360.0
        return image


# 转换图像的色彩空间
class ConvertColor(Module):
    def __init__(self, current="BGR", transform="HSV"):
        super().__init__()
        self.transform = transform
        self.current = current

    def forward(self, image):
        if self.current == "BGR" and self.transform == "HSV":
            image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        elif self.current == "HSV" and self.transform == "BGR":
            image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
        else:
            raise NotImplementedError
        return image


class RandomSaturation(Module):
    def __init__(self, lower=0.5, upper=1.5):
        super().__init__()
        self.lower = lower
        self.upper = upper

    def forward(self, image):
        if random.randint(0, 1):
            # 随机缩放S空间的值
            image[:, :, 1] *= random.uniform(self.lower, self.upper)
        return image


class RandomContrast(Module):
    def __init__(self, lower=0.5, upper=1.5):
        super().__init__()
        self.lower = lower
        self.upper = upper

    def forward(self, image):
        if random.randint(0, 1):
            # 生成随机因子
            alpha = random.uniform(self.lower, self.upper)
            image *= alpha
        return image


class PhotometricDistort(Module):
    def __init__(self):
        super().__init__()
        self.pd = [
            RandomContrast(),  # 随机对比度
            ConvertColor(transform="HSV"),  # 转换色彩空间
            RandomSaturation(),  # 随机饱和度
            # RandomHue(),  # 随机色调
            ConvertColor(current="HSV", transform="BGR"),  # 转换色彩空间
            RandomContrast(),  # 随机对比度
        ]
        self.rand_brightness = RandomBrightness()  # 随机亮度
        self.rand_light_noise = RandomLightingNoise()  # 随机通道交换

    def forward(self, image):
        im = image.copy()
        im = self.rand_brightness(im)
        if random.randint(0, 1):
            distort = Sequential(*self.pd[:-1])
        else:
            distort = Sequential(*self.pd[1:])
        im = distort(im)
        im = self.rand_light_noise(im)
        im = (im - im.min()) / (im.max() - im.min()) * 255
        return im


class RandomMirror(Module):
    def __init__(self):
        super().__init__()

    def forward(self, image, boxes):
        height, width, _ = image.shape
        if random.randint(0, 1):
            # 图像翻转
            image = image[:, ::-1]
            boxes = boxes.copy()
            # 改变标注框
            boxes[:, 0::2] = width - boxes[:, 2::-2]
        if random.randint(0, 1):
            # 图像翻转
            image = image[::-1]
            boxes = boxes.copy()
            # 改变标注框
            boxes[:, 1::2] = height - boxes[:, 3::-2]
        return image, boxes


class Expand(Module):
    def __init__(self, mean=0):
        super().__init__()
        self.mean = mean

    def forward(self, image, boxes):
        if random.randint(0, 1):
            return image, boxes
        # 获取图像的各个维度
        height, width, depth = image.shape
        # 随机缩放尺度
        ratio = random.uniform(1, 4)
        left = random.uniform(0, width * ratio - width)
        top = random.uniform(0, height * ratio - height)
        # 确定缩放后的图像的维度
        expand_image = np.zeros(
            (int(height * ratio), int(width * ratio), depth), dtype=image.dtype
        )
        expand_image[:, :, :] = self.mean
        expand_image[int(top) : int(top + height), int(left) : int(left + width)] = (
            image
        )
        # 返回缩放后的图像
        image = expand_image
        # 将边界框以同等方式缩放
        boxes = boxes.detach().clone()
        boxes += torch.tensor([int(left), int(top), int(left), int(top)]).unsqueeze(0)
        # 返回
        return image, boxes


class RandomSampleCrop(Module):
    def __init__(self):
        super().__init__()
        self.sample_options = (
            None,
            (0.1, None),
            (0.3, None),
            (0.7, None),
            (0.9, None),
            (None, None),
        )

    def forward(self, image, boxes):
        height, width, _ = image.shape
        while True:
            # 随机选择一种裁剪方式
            model = random.choice(self.sample_options)
            # 随机到None直接返回
            if model is None:
                return image, boxes
            # 最大IoU和最小IoU
            min_iou, max_iou = model
            if min_iou is None:
                min_iou = float("-inf")
            if max_iou is None:
                max_iou = float("inf")
            # 迭代50次
            for _ in range(50):
                current_image = image
                # 宽和高随机采样
                w = random.uniform(0.3 * width, width)
                h = random.uniform(0.3 * height, height)
                # 宽高比例不当
                if h / w < 0.5 or h / w > 2:
                    continue
                left = random.uniform(0, width - w)
                top = random.uniform(0, height - h)
                # 框坐标x1,y1,x2,y2
                rect = torch.tensor([int(left), int(top), int(left + w), int(top + h)])
                # 求iou
                overlap = iou(boxes, rect.unsqueeze(0))
                if overlap.min() < min_iou and max_iou < overlap.max():
                    continue
                # 裁剪图像
                current_image = current_image[rect[1] : rect[3], rect[0] : rect[2], :]
                # 中心点坐标
                centers = (boxes[:, :2] + boxes[:, 2:]) / 2.0
                m1 = (rect[0] < centers[:, 0]) * (rect[1] < centers[:, 1])
                m2 = (rect[2] > centers[:, 0]) * (rect[3] > centers[:, 1])
                # 当m1和m2均为正时才保留
                mask = m1 * m2
                if not mask.any():
                    continue
                current_boxes = boxes[mask, :].detach().clone()
                # 根据图像变换调整box
                current_boxes[:, :2] = np.maximum(current_boxes[:, :2], rect[:2])
                current_boxes[:, :2] -= rect[:2]
                current_boxes[:, 2:] = np.minimum(current_boxes[:, 2:], rect[2:])
                current_boxes[:, 2:] -= rect[:2]
                # 返回变换后的图像、box和label
                # 将裁剪后的图像放缩回原图大小
                current_image = cv2.resize(current_image, (width, height))
                # 将边界框放缩回原图大小
                current_boxes = (
                    current_boxes
                    * torch.tensor([width, height, width, height])
                    / torch.tensor([w, h, w, h])
                )
                return current_image, current_boxes
